import os
import shutil
import sqlite3
import argparse
from pathlib import Path

# 定义视频扩展名集合
video_extensions = {'.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.mts', '.mpeg', '.rm', '.webm', '.3gp', '.mp3', '.wav', '.f4v', '.m4a', '.mpg', '.rmvb', '.vob'}

def get_video_info(db_path, video_hash):
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute("SELECT file_path, file_name FROM video_hash WHERE file_hash=?", (video_hash,))
    return c.fetchone()

def main():
    parser = argparse.ArgumentParser(description='Restore video files based on hash values.')
    parser.add_argument('--hash', required=True, help='Directory containing text files with video hashes')
    parser.add_argument('--baseline', required=True, help='Root directory for baseline video library and SQLite database')
    parser.add_argument('--restore', required=True, help='Target directory to restore video files')
    args = parser.parse_args()

    db_path = os.path.join(args.baseline, 'file_hash.db')

    for root, dirs, files in os.walk(args.hash):
        for file in files:
            if file.endswith('.txt'):
                txt_file_path = os.path.join(root, file)
                base_name = file[:-4]  # Remove .txt extension

                # 检查base_name是否以视频扩展名结尾
                if not any(base_name.lower().endswith(ext) for ext in video_extensions):
                    continue  # 跳过非视频文件

                with open(txt_file_path, 'r') as f:
                    video_hash = f.read().strip()

                result = get_video_info(db_path, video_hash)
                if not result:  # 如果没有找到对应的记录
                    print(f"No record found for hash {video_hash} in {txt_file_path}")
                    continue

                source_file_path, source_file_name = result

                # 构建源文件和目标文件的完整路径
                source_full_path = os.path.join(args.baseline, source_file_path)
                target_dir = os.path.join(args.restore, os.path.relpath(root, args.hash))
                os.makedirs(target_dir, exist_ok=True)
                target_full_path = os.path.join(target_dir, base_name)

                try:
                    shutil.copy2(source_full_path, target_full_path)
                    # print(f"Copied {source_full_path} to {target_full_path}")
                except FileNotFoundError:
                    print(f"Source file not found: {source_full_path}")
                except Exception as e:
                    print(f"Error copying file {source_full_path} to {target_full_path}: {e}")

if __name__ == "__main__":
    main()